#ifndef SRC_MPI_CIPHER_MPI_BASE
#define SRC_MPI_CIPHER_MPI_BASE

#include <NTL/mat_ZZ.h>

#include "crypto_helper.hpp"
#include "hypercube/hypercube_hmm_engine.hpp"
#include "log_helper.hpp"
#include "mpi/mpi_util.hpp"

template <typename Scheme, typename ExecOutput> class cipher_mpi_base {
protected:
  static constexpr int MPI_COMM_EXEC_ROOT = 0;

  struct metric {
    // * setup_time
    //   * find_params_time   <- hypercube_hmmcs
    //   * build_context_time <- hypercube_hmmcs
    //   * key_gen_time       <- hmm_crypto_client
    //   * bcast_pk_time      (Broadcast the public key)
    //   * hmm_encoding_time
    double setup_time = 0.0;
    double bcast_pk_time = 0.0;
    double hmm_encoding_time = 0.0;

    // * exec_time
    double exec_time = 0.0;

    // * communication
    double pk_size = 0.0;
    double computing_comm_size = 0.0;

    // * noise bound
    double final_noise = 0.0;
    double init_noise = 0.0;
    double cgemm_noise = 0.0;
    double add_partial_noise = 0.0;

    virtual void clear() {
      exec_time = 0.0;
      computing_comm_size = 0.0;
      final_noise = 0.0;
      init_noise = 0.0;
      cgemm_noise = 0.0;
      add_partial_noise = 0.0;
    }
  };

protected:
  cipher_mpi_base() = delete;
  cipher_mpi_base(int root, MPI_Comm comm)
      : cipher_root_(root), comm_world_(comm) {}
  virtual ~cipher_mpi_base() {
    delete hmm_engine_;
    delete hmm_cagent_;
    delete pk_;
  }

  /**
   * @brief Set the `comm_exec`
   *        Create an exectuion communicator whose rank 0
   *        is rank `cipher_root` in `comm`
   * @param exec_np Size of `comm_exec`
   * @param comm    Old communicator
   */
  void create_comm_exec(int exec_np, MPI_Comm comm);

  /**
   * @brief  broadcast public key in `comm_exec`
   *
   * @return double: Size of the public key
   */
  double bcast_pubkey();

  /**
   * @brief Distribute initial data from `cipher_root` in `comm_world_`
   *
   * @param A
   * @param B
   */
  virtual void init_data(const int *A, const int *B) = 0;
  virtual void exec(ExecOutput &partial_C, bool fhe4d = false) = 0;
  virtual void collect_data(const ExecOutput &partial_C,
                            std::vector<NTL::mat_ZZ> &result,
                            bool local = false) = 0;
  virtual hmm_status<shmm_engine> *hmme_status() = 0;
  virtual void report(std::ostream &out) const = 0;
  virtual void clear() {}

public:
  bool exec_active() const { return active_exec_; }

  hmm_crypto_agent *hmme() { return hmm_cagent_; }

protected:
  MPI_Comm comm_world_;
  int cipher_root_; // in `comm_world_`

  bool active_exec_;
  MPI_Comm comm_exec_;
  int rank_exec_;

  MATRIX_TYPE type_ = MATRIX_TYPE::UNDEF_MATRIX_TYPE;
  shmm_engine *hmm_engine_ = nullptr;

  helib::PubKey *pk_ = nullptr;
  hmm_crypto_agent *hmm_cagent_ = nullptr;

  ExecOutput local_A_;
  ExecOutput local_B_;
};

template <typename Scheme, typename ExecOutput>
double cipher_mpi_base<Scheme, ExecOutput>::bcast_pubkey() {
  std::string pk_str;
  std::stringstream pk_ss;
  if (rank_exec_ == MPI_COMM_EXEC_ROOT) {
    // Cast to `hypercube_hmmcc` to retrieve the public key
    //! Because `hmm_crypto_client` is only an interface,
    //! we must use `dynamic_cast`
    auto hmmcc = dynamic_cast<hmm_crypto_client<Scheme> *>(hmm_cagent_);
    ASSERT_PRINTF(hmmcc != nullptr, "Invalid root\n");
    // Allocate a new public key
    pk_ = new helib::PubKey(hmmcc->public_key());
    pk_->writeTo(pk_ss);
    pk_str = pk_ss.str();
  }
  MPI_Bcast_String(pk_str, MPI_COMM_EXEC_ROOT, comm_exec_);
  if (rank_exec_ != MPI_COMM_EXEC_ROOT) {
    auto hmmcs = static_cast<hmm_crypto_agent *>(hmm_cagent_);
    pk_ss << pk_str;
    pk_ = new helib::PubKey(helib::PubKey::readFrom(pk_ss, *hmmcs->context()));
  }
  return pk_str.size() / 1024.0 / 1024.0;
}

template <typename Scheme, typename ExecOutput>
void cipher_mpi_base<Scheme, ExecOutput>::create_comm_exec(int exec_np,
                                                           MPI_Comm comm) {
  int rank, np;
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &np);

  // sanity check
  ASSERT_PRINTF(exec_np <= np, "Invalid exec_np=%d > |comm|=%d\n", exec_np, np);
  ASSERT_PRINTF(MPI_COMM_EXEC_ROOT < exec_np,
                "Invalid MPI_COMM_EXEC_ROOT=%d > |comm_exec|=%d\n",
                MPI_COMM_EXEC_ROOT, exec_np);

  int rank_exec = (rank - cipher_root_ + MPI_COMM_EXEC_ROOT + np) % np;
  MPI_Comm_split(comm, rank_exec < exec_np ? 0 : MPI_UNDEFINED, rank,
                 &comm_exec_);
  if (comm_exec_ != MPI_COMM_NULL) {
    active_exec_ = true;
    MPI_Comm_rank(comm_exec_, &rank_exec_);
  }
}

double bcast_pubkey(const helib::Context &context, helib::PubKey *&pk, int root,
                    MPI_Comm comm);
void bcast_seckey(const helib::Context &context, helib::SecKey *&sk, int root,
                  MPI_Comm comm);

void ctxt_reduce_add(const helib::Ctxt &send_ctxt, helib::Ctxt *recv_ctxt,
                     int root, MPI_Comm comm);
void ctxt_allreduce_add(const helib::Ctxt &send_ctxt, helib::Ctxt *recv_ctxt,
                        MPI_Comm comm);

void ctxt_ireduce_add(const helib::Ctxt &send_ctxt, helib::Ctxt *recv_ctxt,
                      int root, MPI_Comm comm);

#endif /* SRC_MPI_CIPHER_MPI_BASE */
